#!/usr/bin/env python3
import asyncio
import time
from logger import Logger
import telnetlib3
from device import *

class Connection:
    def __init__(self, info: DeviceInfo, logger, client_id: str):
        self.info = info
        self.logger = logger
        self.reader = None
        self.writer = None
        self.is_connected = False
        self.connect_time = None
        self.last_alive_time = None
        self.session = None
        self.client_id = client_id
        self.lock = asyncio.Lock()
        self.is_sonic = False
        self.last_line = "" # 暂存最后一行输出，用于模式识别

    async def connect(self):
        try:
            if self.is_connected:
                return 1, ""

            # 使用锁保护连接过程，避免多任务并发连接
            async with self.lock:
                # 二次检查连接状态，因为可能在获取锁期间状态已改变
                if self.is_connected:
                    return 1, ""

                # 使用asyncio.wait_for设置超时，因为telnetlib3.open_connection不直接支持timeout参数
                self.reader, self.writer = await asyncio.wait_for(
                    telnetlib3.open_connection(self.info.ip, self.info.port), 
                    timeout=10
                )
                # 读取连接后的首次输出,保留最后一行
                data = await asyncio.wait_for(self.reader.read(4096), timeout=0.5)
                if data:
                    lines = data.splitlines()
                    if lines:
                        self.last_line = lines[-1]
                    else:
                        self.last_line = ""

                self.is_connected = True
                self.connect_time = time.time()
                self.last_alive_time = time.time()
            
            # 注册日志文件
            self.logger.register_connection(self.info.name, self.info.ip, self.info.port)
            
            log_msg = f'成功连接到 {self.info.ip}:{self.info.port}'
            self.logger.log(self.info.name, log_msg)

            return 1, ""
        except Exception as e:
            log_msg = f'连接失败, 设备: {self.info.name}, 请检查网络和配置, 错误: {str(e)}'
            self.logger.server_log(log_msg)
            self.is_connected = False
            return 0, log_msg

    async def send_command(self, command):
        MAX_RECONNECT_ATTEMPTS = 3

        for attempt in range(MAX_RECONNECT_ATTEMPTS):
            try:
                # 1. 连接检查（在循环内进行初始检查）
                if not self.is_connected or self.reader is None or self.writer is None:
                    ret, retstr = await self.connect()
                    if not ret:
                        if attempt == MAX_RECONNECT_ATTEMPTS - 1: # 最后一次尝试失败直接返回
                            return 0, retstr
                        continue  # 重新尝试连接
                    else:
                        mode = get_command_mode(command)
                        ret, ret_str = await self.switch_mode(mode)
                        if not ret:
                            return {
                                    'success': False,
                                    'message': ret_str
                                    }

                # 2. 使用锁保护命令发送和接收过程
                async with self.lock:
                    # 3. 再次检查连接状态（获取锁后）
                    if not self.is_connected or self.reader is None or self.writer is None:
                        if attempt == MAX_RECONNECT_ATTEMPTS - 1:
                            return 0, "连接已断开，无法发送命令"
                        # 尝试重新连接
                        continue

                    # 4. 发送命令
                    try:
                        self.writer.write(command + '\n')
                        await self.writer.drain()
                        await asyncio.sleep(0.5)  # 等待命令执行
                    except Exception as e:
                        if attempt == MAX_RECONNECT_ATTEMPTS - 1:
                            self.is_connected = False
                            return 0, f"发送命令时发生错误: {str(e)}, command: {repr(command)}"
                        self.is_connected = False
                        await asyncio.sleep(1)  # 等待一段时间再重连
                        continue  # 重新发送命令

                    # 5. 读取输出
                    output = []
                    try:
                        data = await asyncio.wait_for(self.reader.read(4096), timeout=0.5)
                        if data:
                            output.append(data)
                    except Exception as e:
                        if attempt == MAX_RECONNECT_ATTEMPTS - 1:
                            return 0, f"读取输出时发生错误: {str(e)}, command: {repr(command)}"

                    # 6. 处理输出数据
                    output_str = ''
                    if output:
                        # 处理输出...
                        output_str = ''.join(output)
                        # 剔除输入的命令本身...
                        output_str = output_str.lstrip('\r\n')

                        if not output_str:
                            # 如果没有输出，尝试再次读取
                            try:
                                await asyncio.sleep(0.5)
                                data = await asyncio.wait_for(self.reader.read(4096), timeout=1.0)
                                if data:
                                    output_str = data.strip()
                            except asyncio.TimeoutError:
                                self.is_connected = False
                                log_msg = f'连接中断，读取输出时超时: {str(e)}, command: {repr(command)}'
                                self.logger.server_log(log_msg)
                                return 0, log_msg
                    else:
                        # 如果没有读取到任何数据，连接可能中断
                        if attempt == MAX_RECONNECT_ATTEMPTS - 1:
                            self.is_connected = False
                            return 0, "[重连失败]"
                        self.is_connected = False
                        await asyncio.sleep(1)
                        continue  # 重新发送命令

                    # 7. 如果执行到这里，说明成功获取输出，记录并返回
                    self.logger.log(self.info.name, f'发送命令: {repr(command)}')
                    self.logger.log(self.info.name, f'命令输出: {output_str}')
                    self.last_alive_time = time.time()
                    lines = output_str.splitlines()
                    self.last_line = lines[-1] if lines else ""
                    return 1, output_str

            except Exception as e:
                error_msg = f'发送命令失败: {str(e)}'
                self.logger.log(self.info.name, error_msg)
                return 0, error_msg

        # 如果循环结束还没返回，说明重试次数用完
        return 0, "重试次数已达上限"

    async def close(self):
        try:
            if self.is_connected and self.writer is not None:
                self.writer.close()
                # 处理wait_closed方法可能不存在的情况
                try:
                    await self.writer.wait_closed()
                except AttributeError:
                    # 如果没有wait_closed方法，等待一小段时间
                    await asyncio.sleep(0.5)
                self.is_connected = False
                log_msg = f'连接已关闭'
                self.logger.log(self.info.name, log_msg)
                
                # 取消注册日志文件
                self.logger.unregister_connection(self.info.name)
                
                return True
            return False
        except Exception as e:
            log_msg = f'关闭连接失败: {str(e)}'
            self.logger.log(self.info.name, log_msg)
            return False

    def is_alive(self):
        return self.is_connected

    def get_info(self):
        return {
            'name': self.info.name,
            'ip': self.info.ip,
            'port': self.info.port,
            'status': '活跃' if self.is_connected else '已关闭',
            'connect_time': self.connect_time
        }
    
    async def handle_login(self):
        """
        处理登录过程（用户名和密码）
        """
        try_count = 0
        max_attempts = 5  # 总尝试次数限制，防止无限循环

        # 发送回车，确认需要输入的信息
        ret, response_str = await self.send_command(f'\n')
        if not ret:
            self.logger.log(self.info.name, f'发送回车失败: {response_str}')
            return ret, response_str

        # 发送用户名

        while True:
            # 检查是否超出总尝试次数
            if try_count >= max_attempts:
                self.logger.log(self.info.name, f'登录尝试次数过多({try_count})，已达上限')
                return 0, f'设备 {self.info.name} 登录尝试次数过多，已达上限'
            
            # 检查连接状态
            if not self.is_connected:
                self.logger.log(self.info.name, '连接已断开，无法继续登录过程')
                return 0, f'设备 {self.info.name} 连接已断开'
            
            str_strip = response_str.lower()
            str_strip = str_strip.strip()

            # 检查是否需要继续登录流程
            if not str_strip.endswith('login:') and not str_strip.endswith('password:'):
                break
            self.logger.log(self.info.name, '开始登录')
            # 处理用户名输入
            if str_strip.endswith('login:'):
                user = self.info.username
                if not user:
                    self.logger.log(self.info.name, f'缺少用户名配置')
                    return 0, f'设备 {self.info.name} 缺少用户名配置'
                self.logger.log(self.info.name, f'设备 {self.info.name} 输入用户名{self.info.username}')
                ret, response_str = await self.send_command(f'{user}')
                if not ret:
                    self.logger.log(self.info.name, f'发送用户名失败: {response_str}')
                    return ret, response_str
                try_count += 1

            # 处理密码输入
            elif str_strip.endswith('password:'):
                password = self.info.password
                if not password:
                    self.logger.log(self.info.name, '缺少密码配置')
                    return 0, f'设备 {self.info.name} 缺少密码配置'
                self.logger.log(self.info.name, f'设备 {self.info.name} 输入登录密码')
                ret, response_str = await self.send_command(f'{password}')
                if not ret:
                    self.logger.log(self.info.name, f'发送密码失败: {response_str}')
                    return ret, response_str
                try_count += 1
                if try_count > 3:
                    self.logger.log(self.info.name, '登录密码错误')
                    return 0, f'设备 {self.info.name} 登录密码错误'
        self.logger.log(self.info.name, f'设备 {self.info.name} 登录成功')
        return 1, response_str

    async def handle_enable_password(self):
        """
        处理特权模式密码输入
        """
        max_attempts = 3  # 最大尝试次数，防止无限循环
        attempt_count = 0
        
        # 发送en命令
        ret, response_str = await self.send_command('en')
        if not ret:
            self.logger.log(self.info.name, f'发送en命令失败: {response_str}')
            return ret, response_str

        while True:
            # 检查连接状态
            if not self.is_connected:
                self.logger.log(self.info.name, '连接已断开，无法继续特权模式切换')
                return 0, f'设备 {self.info.name} 连接已断开'
            
            # 检查是否需要继续特权模式密码输入
            if 'password:' not in response_str.lower():
                return 1, ""
                
            # 检查尝试次数
            attempt_count += 1
            if attempt_count > max_attempts:
                self.logger.log(self.info.name, f'特权模式密码尝试次数过多({attempt_count})，已达上限')
                return 0, f'设备 {self.info.name} 特权模式密码尝试次数过多，已达上限'
                
            # 二选一逻辑：优先使用enable_password，如果没有则使用普通密码
            # 如果enable_password为空，则尝试使用login密码作为备选
            enable_password = self.info.enable_password or self.info.password
        
            if not enable_password:
                self.logger.log(self.info.name, '缺少ENABLE密码配置')
                return 0, f'设备 {self.info.name} 缺少ENABLE密码配置'
                
            self.logger.log(self.info.name, f'设备 {self.info.name} 输入ENABLE密码')
            ret, response_str = await self.send_command(f'{enable_password}')
            if not ret:
                self.logger.log(self.info.name, f'发送特权模式密码失败: {response_str}')
                return ret, response_str
            
        return 1, response_str

    async def get_mode(self):
        # # 发送回车获取当前提示符
        # ret, prompt = await self.send_command('')
        # if not ret:
        #     return ShellMode.INVAILD
        # 去除前后空格
        prompt = self.last_line.strip()
        # 识别当前模式
        current_mode = ShellMode.OTHER
        if isinstance(prompt, str):
            if 'login:' in prompt.lower() or 'password:' in prompt.lower():
                current_mode = ShellMode.NO_LOGIN
            elif prompt.endswith('>'):
                current_mode = ShellMode.NORMAL
            elif prompt.endswith('#') and '(' not in prompt:
                current_mode = ShellMode.ENABLE
            elif prompt.endswith(')#') and '(config' in prompt:
                current_mode = ShellMode.CONFIG
            elif prompt.endswith('~$'):
                self.is_sonic = True
                current_mode = ShellMode.LINUX
            else:
                current_mode = ShellMode.OTHER
        else:
            self.logger.log(self.info.name, f'警告: 提示符不是字符串类型，而是 {type(prompt)}')
        if current_mode not in [ShellMode.LINUX, ShellMode.NO_LOGIN] and self.is_sonic:
            self.logger.log(self.info.name, f'SONIC设备，不支持模式切换')
            return ShellMode.LINUX
        self.logger.log(self.info.name, f'当前模式: {current_mode.name}, 提示符: {prompt}')
        return current_mode

    async def switch_mode(self, mode: ShellMode):
        """
        根据设备返回结果进行模式切换
        
        参数:
            mode: 目标模式,
        
        模式定义:
            - NO_LOGIN: 设备显示"login:"或"Login:"
            - NORMAL: 提示符以>结尾
            - ENABLE: 提示符以#结尾且不包含括号
            - CONFIG: 提示符以(config)#结尾
            - LINUX: 提示符以~$结尾
            - OTHER: 其他情况
        """
        try:
            # if not self.is_connected:
            #     return 0, f'连接 {self.info.name} 未建立'
            current_mode = await self.get_mode()

            # 无论切换到哪种模式，必须先登录
            if current_mode == ShellMode.NO_LOGIN:
                ret, response_str = await self.handle_login()
                if not ret:
                    return ret, response_str
                # 登陆后再次获取
                current_mode = await self.get_mode()
                self.logger.log(self.info.name, f'设备 {self.info.name} 当前模式{current_mode.name}')
            
            if current_mode == ShellMode.LINUX: # linux shell模式，不再切换
                return 1, ""

            if current_mode == ShellMode.INVAILD:
                return 0, f'设备 {self.info.name} 获取模式失败，请检查网络'
            
            if current_mode == ShellMode.NORMAL:
                ret, response_str = await self.handle_enable_password()
                if not ret:
                    return ret, response_str
                # 登陆后再次获取
                current_mode = await self.get_mode()
            
            # 根据当前模式和目标模式执行相应的切换操作
            if mode == ShellMode.ENABLE:  # 目标是ENABLE模式
                if current_mode in [ShellMode.CONFIG, ShellMode.OTHER]:
                    # 配置模式或其他配置模式返回ENABLE模式
                    ret, response_str = await self.send_command('end')
                    if not ret:
                        return ret, response_str
                    self.logger.log(self.info.name, f'已从 {current_mode.name} 切换到 ENABLE模式')
                    
                # 如果已经是ENABLE模式，无需操作
            elif mode == ShellMode.CONFIG:  # 目标是CONFIG模式
                if current_mode == ShellMode.OTHER:
                    # 配置模式或其他配置模式返回ENABLE模式
                    ret, response_str = await self.send_command('end')
                    if not ret:
                        return ret, response_str
                    self.logger.log(self.info.name, f'已从 {current_mode.name} 切换到 ENABLE模式')
                    current_mode = await self.get_mode()
                if current_mode == ShellMode.ENABLE:
                    # ENABLE模式切换至CONFIG模式
                    ret, response_str = await self.send_command('configure terminal')
                    if not ret:
                        return ret, response_str
                    self.logger.log(self.info.name, '已从 ENABLE模式 切换到 CONFIG模式')
            
            return 1, ""
        except Exception as e:
            error_msg = f'模式切换失败: {str(e)}'
            self.logger.log(self.info.name, error_msg)
            return 0, error_msg
